#!/usr/bin/env python3
"""
最终数据维度修复测试
"""

import torch
from torch.utils.data import DataLoader
from data_utils import WeatherDataset, load_config

def test_final_fix():
    """测试最终修复方案"""
    print("最终数据维度修复测试")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        
        # 创建数据集
        dataset = WeatherDataset(
            station_csv=config['data']['station_csv'],
            himawari_dir=config['data']['himawari_dir'],
            historical_days=config['data']['historical_days'],
            future_hours=config['data']['future_hours']
        )
        
        print(f"数据集样本数量: {len(dataset)}")
        
        # 测试单个样本
        if len(dataset) > 0:
            sample = dataset[0]
            print(f"单个样本维度:")
            print(f"  station_data: {sample['station_data'].shape}")
            print(f"  himawari_data: {sample['himawari_data'].shape}")
            print(f"  future_data: {sample['future_data'].shape}")
        
        # 测试数据加载器 - 使用自定义collate函数
        from train import custom_collate_fn
        
        dataloader = DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=custom_collate_fn)
        
        # 获取一个批次
        batch = next(iter(dataloader))
        
        print(f"\n批次数据维度 (使用自定义collate函数):")
        print(f"  station_data: {batch['station_data'].shape}")
        print(f"  himawari_data: {batch['himawari_data'].shape}")
        print(f"  future_data: {batch['future_data'].shape}")
        
        # 检查维度一致性
        station_shapes = [item['station_data'].shape for item in dataset]
        himawari_shapes = [item['himawari_data'].shape for item in dataset]
        future_shapes = [item['future_data'].shape for item in dataset]
        
        print(f"\n维度统计:")
        print(f"  station_data 形状集合: {set(station_shapes)}")
        print(f"  himawari_data 形状集合: {set(himawari_shapes)}")
        print(f"  future_data 形状集合: {set(future_shapes)}")
        
        # 检查是否所有样本都有相同的维度
        if len(set(station_shapes)) == 1:
            print("✅ station_data 维度一致")
        else:
            print("❌ station_data 维度不一致")
            print(f"  不同维度: {set(station_shapes)}")
            
        if len(set(himawari_shapes)) == 1:
            print("✅ himawari_data 维度一致")
        else:
            print("❌ himawari_data 维度不一致")
            print(f"  不同维度: {set(himawari_shapes)}")
            
        if len(set(future_shapes)) == 1:
            print("✅ future_data 维度一致")
        else:
            print("❌ future_data 维度不一致")
            print(f"  不同维度: {set(future_shapes)}")
            
        # 测试模型兼容性
        from models import CombinedModel
        
        model = CombinedModel(config)
        print(f"\n模型测试:")
        
        # 使用批次数据进行前向传播测试
        with torch.no_grad():
            predictions = model(batch['station_data'], batch['himawari_data'])
            print(f"  模型输出: {predictions.shape}")
            print(f"  期望输出: {batch['future_data'].shape}")
            
            if predictions.shape == batch['future_data'].shape:
                print("✅ 模型输出维度正确")
            else:
                print("❌ 模型输出维度错误")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = test_final_fix()
    if success:
        print("\n🎉 最终修复成功！")
    else:
        print("\n⚠️ 最终修复仍然存在问题")
